#!/usr/bin/python
# coding:utf-8

"""
@author: 信长华
@contact: skywater@gmail.com
@software: PyCharm
@file: middlewares.py
@time: 2021/7/31 22:14
"""
from typing import Optional
from uuid import uuid4

from contextvars import ContextVar, Token
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint

CORRELATION_ID_CTX_KEY = "correlation_id"
REQUEST_ID_CTX_KEY = "request_id"

_correlation_id_ctx_var: ContextVar[Optional[str]] = ContextVar(
    CORRELATION_ID_CTX_KEY, default=None
)
_request_id_ctx_var: ContextVar[Optional[str]] = ContextVar(
    REQUEST_ID_CTX_KEY, default=None
)


def get_correlation_id() -> str:
    return _correlation_id_ctx_var.get()


def get_request_id() -> str:
    return _request_id_ctx_var.get()


# 中间件参考 https://github.com/tiangolo/fastapi/issues/397
# 日志参数动态更新参考 https://www.escapelife.site/posts/d4521b7.html
# 也可以参考这个 https://gist.github.com/ddanier/ead419826ac6c3d75c96f9d89bea9bd0
# 日志参考 # https://www.v2ex.com/t/621041
# 注意这个中间件的加载顺序
class RequestContextLogMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        correlation_id_token: Token = _correlation_id_ctx_var.set(
            request.headers.get("X-Correlation-ID", str(uuid4()))
        )
        request_id_token: Token = _request_id_ctx_var.set(str(uuid4()))
        response = await call_next(request)
        response.headers["X-Correlation-ID"] = get_correlation_id()
        response.headers["X-Request-ID"] = get_request_id()

        _correlation_id_ctx_var.reset(correlation_id_token)
        _request_id_ctx_var.reset(request_id_token)

        return response
